#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <sched.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>


/*
 * set process's CPU affinity mask.
 * systemcall: 
 * int sched_setaffinity(pid_t pid, size_t cpusetsize,
 *                       cpu_set_t *mask);
 */

/*
 * DESCRIPTION
 * A  process's  CPU  affinity mask determines the set of CPUs on which it is eligible to run.  On a multiprocessor system, setting the CPU affinity mask can be used to
 * obtain performance benefits.  For example, by dedicating one CPU to a particular process (i.e., setting the affinity mask of that process to specify  a  single  CPU,
 * and  setting the affinity mask of all other processes to exclude that CPU), it is possible to ensure maximum execution speed for that process.  Restricting a process
 * to run on a single CPU also avoids the performance cost caused by the cache invalidation that occurs when a process ceases to execute on one CPU and then recommences
 * execution on a different CPU.
 * 
 * A  CPU  affinity  mask  is  represented  by  the  cpu_set_t  structure,  a  "CPU set", pointed to by mask.  A set of macros for manipulating CPU sets is described in
 * CPU_SET(3).
 * 
 * sched_setaffinity() sets the CPU affinity mask of the process whose ID is pid to the value specified by mask.  If pid is zero, then the calling process is used.  The
 * argument cpusetsize is the length (in bytes) of the data pointed to by mask.  Normally this argument would be specified as sizeof(cpu_set_t).
 * 
 * If the process specified by pid is not currently running on one of the CPUs specified in mask, then that process is migrated to one of the CPUs specified in mask.
 * 
 * sched_getaffinity()  writes  the affinity mask of the process whose ID is pid into the cpu_set_t structure pointed to by mask.  The cpusetsize argument specifies the
 * size (in bytes) of mask.  If pid is zero, then the mask of the calling process is returned.
 * 
 * RETURN VALUE
 * On success, sched_setaffinity() and sched_getaffinity() return 0.  On error, -1 is returned, and errno is set appropriately.
 * 
 * ERRORS
 * EFAULT A supplied memory address was invalid.
 * 
 * EINVAL The affinity bit mask mask contains no processors that are currently physically on the system and permitted to the process according to any restrictions  that
 * may be imposed by the "cpuset" mechanism described in cpuset(7).
 * 
 * EINVAL (sched_getaffinity() and, in kernels before 2.6.9, sched_setaffinity()) cpusetsize is smaller than the size of the affinity mask used by the kernel.
 * 
 * EPERM  (sched_setaffinity())  The calling process does not have appropriate privileges.  The caller needs an effective user ID equal to the user ID or effective user
 * ID of the process identified by pid, or it must possess the CAP_SYS_NICE capability.
 * 
 * ESRCH  The process whose ID is pid could not be found.
 * 
 * VERSIONS
 * The CPU affinity system calls were introduced in Linux kernel 2.5.8.  The system call wrappers were  introduced  in  glibc  2.3.   Initially,  the  glibc  interfaces
 * included  a cpusetsize argument, typed as unsigned int.  In glibc 2.3.3, the cpusetsize argument was removed, but was then restored in glibc 2.3.4, with type size_t.
 * 
 * CONFORMING TO
 * These system calls are Linux-specific.
 * 
 * 
 * NOTES
 * After a call to sched_setaffinity(), the set of CPUs on which the process will actually run is the intersection of the set specified in the mask argument and the set
 * of  CPUs  actually present on the system.  The system may further restrict the set of CPUs on which the process runs if the "cpuset" mechanism described in cpuset(7)
 * is being used.  These restrictions on the actual set of CPUs on which the process will run are silently imposed by the kernel.
 * 
 * sched_setscheduler(2) has a description of the Linux scheduling scheme.
 * 
 * The affinity mask is actually a per-thread attribute that can be adjusted independently for each of the threads in a thread group.  The value returned from a call to
 * gettid(2)  can  be passed in the argument pid.  Specifying pid as 0 will set the attribute for the calling thread, and passing the value returned from a call to get-
 * 	pid(2) will set the attribute for the main thread of the thread group.  (If you are using the POSIX  threads  API,  then  use  pthread_setaffinity_np(3)  instead  of
 * 			sched_setaffinity().)
 * 
 * 	A child created via fork(2) inherits its parent's CPU affinity mask.  The affinity mask is preserved across an execve(2).
 * 
 * 	This  manual  page  describes  the  glibc interface for the CPU affinity calls.  The actual system call interface is slightly different, with the mask being typed as
 * 	unsigned long *, reflecting the fact that the underlying implementation of CPU sets is a simple bit mask.   On  success,  the  raw  sched_getaffinity()  system  call
 * 	returns the size (in bytes) of the cpumask_t data type that is used internally by the kernel to represent the CPU set bit mask.
 */

/* 
 * Data structure to describe CPU mask.
 * typedef struct
 * {
 * 	  __cpu_mask __bits[__CPU_SETSIZE / __NCPUBITS];
 * } cpu_set_t;
 */

int main(int argc, char * argv[])
{
	if(argc != 3)
	{
		printf("\nUsgae: %s <pid> <cpu_id_list>\n", argv[0]);
		printf("you should get the available cpu number on your system first,then input the mask in format 0+1+2+3\n");
		printf("\nfor example: ./set_affinity 0 0\n");
		printf("for example: ./set_affinity 0 1\n");
		printf("for example: ./set_affinity 0 0+1\n");
		printf("\nthe digit numbers are cpu id numbers, and started from 0,so it all should be less than the total cpu number on your system.\n\n");
		exit(1);
	}

	int pid = atoi(argv[1]);

	cpu_set_t mask;
	CPU_ZERO(&mask); 
	//printf("cpu_set size: %d\n", sizeof(cpu_set_t));

	
	/*
	 * help to check available cpu id.
	 */
	int cpu_nr = sysconf(_SC_NPROCESSORS_CONF);
	printf("system's available cpu number: %d\n", cpu_nr);

	/*
	 * the input cpu id list should at least contain one available cpu id,
	 * other wise,the sched_setaffinity() will return -1.
	 */
	char *s1;
	s1 = strtok(argv[2],"+");
	if(NULL != s1)
	{
		CPU_SET(atoi(s1), &mask);
		printf("cpu ");
		printf("%s", s1);

		while((s1 = strtok(NULL,"+")) != NULL)
		{
			CPU_SET(atoi(s1), &mask);
			printf(",%s", s1);
		}
		printf(" added.\n");
		CPU_SET(2, &mask);
	}

	//printf("CPU_COUNT() of set:    %d\n", CPU_COUNT(&mask));
	int ret = sched_setaffinity(pid, sizeof(cpu_set_t), &mask);
	if(ret == -1)
	{
		switch(errno)
		{
			case EFAULT:
				printf("A supplied memory address was invalid.\n");
				break;
			case EINVAL:
				printf("Invalid argument, the input cpu id list should at least contain one available cpu id.\n");
				break;
			case EPERM:
				printf("The calling process does not have appropriate privileges.\n");
				break;
			case ESRCH:
				printf("The process whose ID is pid could not be found.\n");
				break;
			default:
				printf("unknown error.\n");
				break;
		}
		printf("Result: failed\n");
		exit(1);
	}
	printf("Result: success\n");

	return 0;
}
